#!/usr/bin/env python
import numpy as np
import matplotlib.pylab as plt
from cosat.slew import csst_slew
from argparse import ArgumentParser

tranges = [ [0, 1],
            [1, 2],
            [2, 3],
            [3, 4],
            [4, 5],
            [5, 6],
            [6, 7],
            [7, 8],
            [8, 9],
            [9,11] ]
labels = ['Year-1','Year-2','Year-3','Year-4','Year-5',
            'Year-6','Year-7','Year-8','Year-9','Year-10']

def plotSlewAngeStats(obs_jdates, obs_slew_angles):
    print('COSAT> running slew angle statistics')
    t = (obs_jdates-obs_jdates.min()) / 365.25

    density=False
    histtype='step'
    angs = []
    for i in range(10):
        idx = np.logical_and( t >= tranges[i][0], t < tranges[i][1] )
        angs.append(obs_slew_angles[idx])

    logscale=False
    
    ncols=3
    plt.figure(figsize=(12,9))

    plt.subplot(2,2,1)
    plt.hist(angs, bins=15, range=[0,15], stacked=True, log=logscale, label=labels)
    plt.legend(loc='best', frameon=False, ncols=ncols)
    plt.xlabel(r'$0^{\degree} \leq \Delta\theta < 15^{\degree}$', fontsize=14)
    plt.ylabel('Slew Counts', fontsize=14)

    plt.subplot(2,2,2)
    plt.hist(angs, bins=35, range=[15,50], stacked=True, log=logscale, label=labels)
    plt.legend(loc='best', frameon=False, ncols=ncols)
    plt.xlabel(r'$15^{\degree} \leq \Delta\theta < 50^{\degree}$', fontsize=14)
    plt.ylabel('Slew Counts', fontsize=14)

    plt.subplot(2,2,3)
    plt.hist(angs, bins=50, range=[50,100], stacked=True, log=logscale, label=labels)
    plt.legend(loc='best', frameon=False, ncols=ncols)
    plt.xlabel(r'$50^{\degree} \leq \Delta\theta < 100^{\degree}$', fontsize=14)
    plt.ylabel('Slew Counts', fontsize=14)

    plt.subplot(2,2,4)
    plt.hist(angs, bins=80, range=[100,180], stacked=True, log=logscale, label=labels)
    plt.legend(loc='best', frameon=False, ncols=ncols)
    plt.xlabel(r'$100^{\degree} \leq \Delta\theta < 180^{\degree}$', fontsize=14)
    plt.ylabel('Slew Counts', fontsize=14)

    plt.tight_layout()
    plt.savefig('SlewCountsStats.png')
    plt.show()

def plotSlewTimeStats_v1(obs_jdates, obs_slew_angles, slew_opt=0, slew_time_interp_kind='linear'):
    """
    Deprecated.
    """
    print('COSAT> running slew time statistics')
    t = (obs_jdates-obs_jdates.min()) / 365.25

    slew = csst_slew(slew_opt=slew_opt, kind=slew_time_interp_kind)

    density=False
    histtype='step'
    angs = []
    slew_times = slew.getSlewTime(obs_slew_angles) / 86400 # use days instead of seconds

    tot_slew_time = np.sum(slew_times) # days

    weights = []
    for i in range(10):
        idx = np.logical_and( t >= tranges[i][0], t < tranges[i][1] )
        angs.append(obs_slew_angles[idx])
        weights.append(slew_times[idx])

    logscale=False
    
    ncols=3
    plt.figure(figsize=(12,9))

    plt.subplot(2,2,1)
    plt.hist(angs, bins=15, range=[0,15], stacked=True, weights=weights, log=logscale, label=labels)
    plt.legend(loc='best', frameon=False, ncols=ncols)
    plt.xlabel(r'$0^{\degree} \leq \Delta\theta < 15^{\degree}$', fontsize=14)
    plt.ylabel('Slew Time (days)', fontsize=14)
    idx = np.logical_and( obs_slew_angles>=0, obs_slew_angles < 15)
    fracf = np.sum(slew_times[idx]) / tot_slew_time
    frac = ' {:.2f} days / {:.2f} days [{:>.2f}%]'.format(np.sum(slew_times[idx]), tot_slew_time, fracf*100)
    plt.title(frac)

    plt.subplot(2,2,2)
    plt.hist(angs, bins=35, range=[15,50], stacked=True, weights=weights, log=logscale, label=labels)
    plt.legend(loc='best', frameon=False, ncols=ncols)
    plt.xlabel(r'$15^{\degree} \leq \Delta\theta < 50^{\degree}$', fontsize=14)
    plt.ylabel('Slew Time (days)', fontsize=14)
    idx = np.logical_and( obs_slew_angles>=15, obs_slew_angles < 50)
    fracf = np.sum(slew_times[idx]) / tot_slew_time
    frac = ' {:.2f} days / {:.2f} days [{:>.2f}%]'.format(np.sum(slew_times[idx]), tot_slew_time, fracf*100)
    plt.title(frac)

    plt.subplot(2,2,3)
    plt.hist(angs, bins=50, range=[50,100], stacked=True, weights=weights, log=logscale, label=labels)
    plt.legend(loc='best', frameon=False, ncols=ncols)
    plt.xlabel(r'$50^{\degree} \leq \Delta\theta < 100^{\degree}$', fontsize=14)
    plt.ylabel('Slew Time (days)', fontsize=14)
    idx = np.logical_and( obs_slew_angles>=50, obs_slew_angles < 100)
    fracf = np.sum(slew_times[idx]) / tot_slew_time
    frac = ' {:.2f} days / {:.2f} days [{:>.2f}%]'.format(np.sum(slew_times[idx]), tot_slew_time, fracf*100)
    plt.title(frac)

    plt.subplot(2,2,4)
    plt.hist(angs, bins=80, range=[100,180], stacked=True, weights=weights, log=logscale, label=labels)
    plt.legend(loc='best', frameon=False, ncols=ncols)
    plt.xlabel(r'$100^{\degree} \leq \Delta\theta < 180^{\degree}$', fontsize=14)
    plt.ylabel('Slew Time (days)', fontsize=14)
    idx = np.logical_and( obs_slew_angles>=100, obs_slew_angles < 180)
    fracf = np.sum(slew_times[idx]) / tot_slew_time
    frac = ' {:.2f} days / {:.2f} days [{:>.2f}%]'.format(np.sum(slew_times[idx]), tot_slew_time, fracf*100)
    plt.title(frac)

    plt.tight_layout()
    plt.savefig('SlewTimeStats.png')
    plt.show()

def plotSlewTimeStats_v2(obs_jdates, obs_slew_angles, slew_opt=0, slew_time_interp_kind='linear'):
    print('COSAT> running slew time statistics')
    slew = csst_slew(slew_opt=slew_opt, kind=slew_time_interp_kind)
    slew_times = slew.getSlewTime(obs_slew_angles) / 86400 # use days instead of seconds
    tot_slew_time = np.sum(slew_times) # days

    slew_counts1 = []        # 机动次数统计
    slew_counts2 = []        # 机动次数统计
    slew_time_counts1 = []   # 激动时间统计
    slew_time_counts2 = []   # 激动时间统计

    slew_angle_intervals = np.linspace(0,180,181)
    for i in range(len(slew_angle_intervals)-1):
        idx1 = obs_slew_angles >= slew_angle_intervals[i]
        idx2 = obs_slew_angles <= slew_angle_intervals[i]

        slew_counts1.append( np.sum(idx1) / len(obs_slew_angles) )
        slew_counts2.append( np.sum(idx2) / len(obs_slew_angles) )

        slew_time_counts1.append( np.sum(slew_times[idx1]) / tot_slew_time )
        slew_time_counts2.append( np.sum(slew_times[idx2]) / tot_slew_time )

    fig = plt.figure(figsize=(9,7))
    ax1 = plt.subplot(2,1,1)
    
    h, b,  = np.histogram(obs_slew_angles,bins=180)
    ax1.plot((b[1:]+b[:-1])/2, h/np.max(h), 'k-', label=r'Slew counts ratio (bin size $\Delta\theta = 1\degree$)')
    ax1.plot(slew_angle_intervals[:-1], slew_counts1, 'b--', label=r'Slew counts frac: $\Delta\theta \geq *$')
    ax1.plot(slew_angle_intervals[:-1], slew_counts2, 'r-.', label=r'Slew counts frac: $\Delta\theta < *$')
    ax1.grid(linestyle='dashed')
    ax1.set_xlim(-5,180)
    ax1.set_xlabel(r'$\Delta\theta$', fontsize=14)
    ax1.legend(loc='best',frameon=False, fontsize=14)
    ax1.set_title('Total # of slews: {}'.format(len(obs_slew_angles)))

    ax2 = plt.subplot(2,1,2)
    ax2.plot(slew_angle_intervals[:-1], slew_time_counts1, 'b--', label=r'Slew time frac: $\Delta\theta \geq *$')
    ax2.plot(slew_angle_intervals[:-1], slew_time_counts2, 'r-.', label=r'Slew time frac: $\Delta\theta < *$')
    ax2.grid(linestyle='dashed')
    ax2.set_xlim(-5,180)
    ax2.set_xlabel(r'$\Delta\theta$', fontsize=14)
    ax2.legend(loc='best',frameon=False, fontsize=14)
    ax2.set_title('Total slew time: {:.2f} days'.format(tot_slew_time))

    plt.tight_layout()
    plt.savefig('SlewStats.png')
    plt.show()
    

def runSlewStats():
    parser = ArgumentParser()
    parser.add_argument('obs_seq', type=str)
    parser.add_argument('--slew_time','-t',action='store_true',
        help='If true, then make statistic of the slew time instead of angles.')
    args = parser.parse_args()

    d = np.loadtxt(args.obs_seq)

    jd = d[:,0]     # Julian dates
    ang = d[:,17]   # slew angles

    if args.slew_time is False:
        plotSlewAngeStats(jd, ang)
    else:
        plotSlewTimeStats_v2(jd, ang)


if __name__ == '__main__':
    runSlewStats()